/*
 * Description: process preloader
 */
#include "process_preloader.h"
#include "Windows.h"
#include "Psapi.h"
#include <algorithm>
#include <chrono>
#include <string>
#include <sstream>
#include <thread>
#include <vector>
#include "common/guest_utils.h"
#include "log.h"
#include "performance/preload/preload_app_path_manager.h"
#include "performance/utils/perf_utils.h"

namespace PERF {

using namespace COMMON;

constexpr const int DEFAULT_PRELOAD_DELAY_SECOND = 30;
const std::string PRELOAD_HELPER_PATH = "C:\\ProgramData\\performance\\PreloadHelper.exe";
const std::string OP_PRELOAD = "preload";
const std::string OP_AWAKEN = "awaken";

typedef LONG(NTAPI *NtSuspendProcess)(IN HANDLE ProcessHandle);
typedef LONG(NTAPI *NtResumeProcess)(IN HANDLE ProcessHandle);

void ProcessPreloader::SuspendProcess(const DWORD pid)
{
    if (pid == 0) {
        return;
    }
    HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pid);
    if (hProcess != INVALID_HANDLE_VALUE) {
        EmptyWorkingSet(hProcess);
        NtSuspendProcess fNtSuspendProcess =
            (NtSuspendProcess)GetProcAddress(GetModuleHandle("ntdll"), "NtSuspendProcess");
        if (fNtSuspendProcess != nullptr) {
            fNtSuspendProcess(hProcess);
        }
        CloseHandle(hProcess);
    }
}

void ProcessPreloader::ResumeProcess(const DWORD pid)
{
    if (pid == 0) {
        return;
    }
    HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pid);
    if (hProcess == INVALID_HANDLE_VALUE) {
        return;
    }
    NtResumeProcess fNtResumeProcess = (NtResumeProcess)GetProcAddress(GetModuleHandle("ntdll"), "NtResumeProcess");
    if (fNtResumeProcess != nullptr) {
        fNtResumeProcess(hProcess);
    }
    CloseHandle(hProcess);
}

DWORD SingleProcessPreloader::GetPid()
{
    return pid;
}

void SingleProcessPreloader::DoPreload()
{
    std::thread tPreload([&]() {
        PreloadAppPathManager manager;
        std::string path;
        if (!manager.GetPath(GetProcessName(), path)) {
            return;
        }
        std::stringstream cmdStream;
        cmdStream << PRELOAD_HELPER_PATH << " " << OP_PRELOAD << " " << GetProcessName() << " \"" << path << "\"";
        std::string preloadHelperCmd = cmdStream.str();
        LOG_INFO("[PERF - DoPreload] - executing cmd: [{}]", preloadHelperCmd);
        GuestUtils::ExecCmdWithExitCode(preloadHelperCmd, true, pid);
        LOG_INFO("[PERF - DoPreload] - pname: {}, pid: {}", GetProcessName(), pid);
        std::this_thread::sleep_for(std::chrono::seconds(DEFAULT_PRELOAD_DELAY_SECOND));
        Suspend();
    });
    tPreload.detach();
}

void SingleProcessPreloader::DoAwaken(const std::string &filepath)
{
    std::thread tAwaken([this, filepath]() {
        Resume();
        std::stringstream cmdStream;
        cmdStream << PRELOAD_HELPER_PATH << " " << OP_AWAKEN << " " << GetProcessName()
                  << (filepath == "" ? "" : " \"" + filepath + "\"");
        std::string preloadHelperCmd = cmdStream.str();
        DWORD exitCode = 0;
        LOG_INFO("[PERF - DoAwaken] - executing cmd: [{}]", preloadHelperCmd);
        GuestUtils::ExecCmdWithExitCode(preloadHelperCmd, true, exitCode);
        LOG_INFO("[PERF - DoAwaken] - awakening {} result: {}", GetProcessName(), exitCode);
    });
    tAwaken.detach();
}

void SingleProcessPreloader::Suspend()
{
    if (pid != 0) {
        LOG_INFO("[PERF - Suspend] - suspending process: {}", pid);
        SuspendProcess(pid);
    }
}

void SingleProcessPreloader::Resume()
{
    if (pid != 0) {
        LOG_INFO("[PERF - Resume] - Resuming process: {}", pid);
        ResumeProcess(pid);
    }
    pid = 0;
}

std::vector<DWORD> MultiProcessPreloader::GetPids()
{
    return pids;
}

void MultiProcessPreloader::Suspend()
{
    for (DWORD pid : pids) {
        if (pid != 0) {
            SuspendProcess(pid);
        }
    }
}

void MultiProcessPreloader::Resume()
{
    for (DWORD pid : pids) {
        if (pid != 0) {
            ResumeProcess(pid);
        }
    }
    pids.clear();
}

}  // namespace PERF